

import cv2
import numpy as np

import glob
import json
#coding : utf-8
import os
import abc
import xml.dom.minidom as xml


class_name = ["truck_bed_cover", "dirt",
              "empty_bed", "back_plate", "side_plate"]
categoriesmask = [
    {
        "supercategory": "none",
        "id": 0,
        "name": "truck_bed_cover"
    },
    {
        "supercategory": "none",
        "id": 1,
        "name": "dirt"
    },
    {
        "supercategory": "none",
        "id": 2,
        "name": "empty_bed"
    },
    {
        "supercategory": "none",
        "id": 3,
        "name": "back_plate"
    },
    {
        "supercategory": "none",
        "id": 4,
        "name": "side_plate"
    }
]


class XmlReader(object):
    __metaclass__ = abc.ABCMeta

    def __init__(self):
        pass

    def read_content(self, filename):
        content = None
        if (False == os.path.exists(filename)):
            return content
        filehandle = None
        try:
            filehandle = open(filename, 'rb')
        except FileNotFoundError as e:
            print(e.strerror)
        try:
            content = filehandle.read()
        except IOError as e:
            print(e.strerror)
        if (None != filehandle):
            filehandle.close()
        if(None != content):
            return content.decode("utf-8", "ignore")
        return content

    @abc.abstractmethod
    def load(self, filename):
        pass


ROOT_PATH ="/home/data"  #os.getcwd()


def put_mask(image, point):

    zeros = np.zeros((image.shape), dtype=np.uint8)
    # zeros_mask = cv2.rectangle(zeros, (bbox[0], bbox[1]), (bbox[2], bbox[3]),
    #                color=(0,0,255), thickness=-1 ) #thickness=-1 表示矩形框内颜色填充
    cv2.fillPoly(zeros, [point], (0, 0, 255), 1)
    gamma = 0

    mask_img = cv2.addWeighted(image, 1, zeros, 0.9, gamma)
    return mask_img


class XmlTester(XmlReader):
    def __init__(self):
        XmlReader.__init__(self)

    def load(self, filename):
        filecontent = XmlReader.read_content(self, filename)
        if None != filecontent:
            # print(filename)
            dom = xml.parseString(filecontent)
            im_size = dom.getElementsByTagName('size')[0]

            im_w = int((im_size.getElementsByTagName(
                'width')[0]).childNodes[0].data)
            im_h = int((im_size.getElementsByTagName(
                "height")[0]).childNodes[0].data)
            self.im_shape = np.array([im_w, im_h])
            self.bbox = []
            self.name = dom.getElementsByTagName(
                'filename')[0].childNodes[0].data
            self.obj_name = []
            for obj in dom.getElementsByTagName('object'):
                box = obj.getElementsByTagName('bndbox')[0]

                b_xmin = int((box.getElementsByTagName(
                    "xmin")[0]).childNodes[0].data)
                b_ymin = int((box.getElementsByTagName(
                    "ymin")[0]).childNodes[0].data)
                b_xmax = int((box.getElementsByTagName(
                    "xmax")[0]).childNodes[0].data)
                b_ymax = int((box.getElementsByTagName(
                    "ymax")[0]).childNodes[0].data)
                b_name = obj.getElementsByTagName(
                    "name")[0].childNodes[0].data
                self.obj_name.append(b_name)
                # print(self.name)
                self.bbox.append([[b_xmin, b_ymin], [b_xmax, b_ymax]])
            self.bbox = np.array(self.bbox)
            self.polygon = []
            for obj in dom.getElementsByTagName('polygon'):
                o_cls = obj.getElementsByTagName('class')[0].childNodes[0].data
                box = obj.getElementsByTagName(
                    'points')[0].childNodes[0].data.replace(";", ",").split(",")
                if o_cls != "back_plate" or  o_cls !="side_plate":
                    self.polygon.append(
                        [o_cls, np.array(box).astype(np.int32).reshape(-1, 2)])
            # print(self.polygon)
                # b_xmin = int((box.getElementsByTagName(
                #     "xmin")[0]).childNodes[0].data)
                # b_ymin = int((box.getElementsByTagName(
                #     "ymin")[0]).childNodes[0].data)
                # b_xmax = int((box.getElementsByTagName(
                #     "xmax")[0]).childNodes[0].data)
                # b_ymax = int((box.getElementsByTagName(
                #     "ymax")[0]).childNodes[0].data)
                # self.name = dom.getElementsByTagName(
                #     'filename')[0].childNodes[0].data
                # print(self.name)
                # self.bbox.append([[b_xmin, b_ymin], [b_xmax, b_ymax]])
            # print(bbox)

            # print(self.polygon)
            return self.im_shape, self.bbox

    def xyxy2yolo(self):
        self.bboxyolo = []
        for bx in self.bbox:
            cent = (bx[1]+bx[0])/2
            dwh = bx[1]-bx[0]
            # self.bboxyolo.append([bx[0]/self.im_shape,dwh/self.im_shape])
            self.bboxyolo.append([cent/self.im_shape, dwh/self.im_shape])

        # for o_cls, pt in self.polygon:
        #     # print(o_cls[-5:])
        #     if o_cls[-5:]=="plate":
        #         # print(sorted(pt))
        #         plx=pt.min(0)
        #         prx=pt.max(0)
        #         self.bboxyolo.append([plx/self.im_shape,prx/self.im_shape])
        return np.array(self.bboxyolo).reshape(-1, 4), self.obj_name

    def imshow(self):
        img = cv2.imread(ROOT_PATH+"/"+self.name)
        for bx in self.bbox:
            cv2.rectangle(img, bx[0], bx[1], (0, 255, 0), 1, 1)
        # cv2.imshow("a", img)
        # cv2.waitKey()
        # img=cv2.cvtColor(img, cv2.COLOR_BGR2BGRA);
        for clas, point in self.polygon:
            # cv2.polylines(img, [point], 1, (0,0,0,255))
            # print(point)
            img = put_mask(img, point)
        size_decrease = (int(img.shape[1]/2), int(img.shape[0]/2))
        img_decrease = cv2.resize(
            img, size_decrease, interpolation=cv2.INTER_CUBIC)
        # mask = im
        # cv2.imshow("Mask", mask)
        # masked = cv2.bitwise_and(image, image, mask=mask)
        cv2.imshow("Mask to Image", img_decrease)
        cv2.waitKey(0)


if __name__ == "__main__":

    xmls = glob.glob("%s/1441/*.xml" % ROOT_PATH)
    os.system("rm -r  %s/yolo3d"%ROOT_PATH)
    os.system("mkdir -p %s/yolo3d/labels"%ROOT_PATH)
    os.system("mkdir -p %s/yolo3d/images"%ROOT_PATH)
    ftxt = open("%s/yolo3d/train.txt"%ROOT_PATH, mode="w")

    annotations_list, image_list = [], []
    img_cnt, anno_cnt = 0, 0
    for i, xml_path in enumerate(xmls):
        reader = XmlTester()
        im_shape, bbox = reader.load(xml_path)

        if len(reader.polygon) > 0:

            img_path = xml_path[:-3]+"jpg"
            img_src_size=im_shape
            wpr="%s/yolo3d/images/%d.jpg\n" % (ROOT_PATH, img_cnt)
            os.system(f"ln -s {img_path} {wpr[:-1]}")
            lb_img_path=img_path[:-3]+"png"
            # os.system(f"ln -s {lb_img_path} {ROOT_PATH}/yolo3d/labels/{img_cnt}.png")
            ftxt.write(wpr)
            mask=cv2.imread(lb_img_path)
            cv2.imwrite(f"{ROOT_PATH}/yolo3d/labels/{img_cnt}.png",mask)
            # flab=open(f"{ROOT_PATH}/yolo3d/labels/{img_cnt}.txt",mode="w")
            # print(reader.polygon)
            for cls, box in reader.polygon:
                fb = box.reshape(-1, 2)/np.array([img_src_size[1],img_src_size[0]])
                stra=np.str_(fb.reshape(-1).flatten())
                # flab.write(f"{class_name.index(cls)} "+ stra[1:-1].replace('\n'," ")+"\n")
            img_cnt+=1
            # flab.close()
        # else:
            # print(i, xml_path)
    print("xmls:",len(xmls))
    ftxt.close()
